#!/usr/bin/env python3
# coding: utf-8
# File: answer_search.py
# Author: lhy<lhy_in_blcu@126.com,https://huangyong.github.io>
# Date: 18-10-5

from py2neo import Graph, Path
import re
from operator import itemgetter
from itertools import groupby


class AnswerSearcher:
    def __init__(self):
        self.g = Graph("http://localhost:7474", auth=("neo4j", "123456"))
        self.num_limit = 20
        self.count = 0
        # 映射字典
        self.categories = {
            'Disease': 0,
            'Check': 1,
            'Symptom': 2,
            'Producer': 3,
            'Drug': 4,
            'Food': 5,
            'Department': 6
        }

    '''执行cypher查询，并返回相应结果'''

    def search_main(self, sqls):
        # 计数器，值为结点id
        self.count = 0

        nodes = []
        edges = []
        entities = []
        relations = []
        answer = []
        for sql_ in sqls:
            question_type = sql_['question_type']
            relations.append(question_type)
            entities.extend(sql_['entities'])
            queries = sql_['sql']
            answers = []

            for query in queries:
                ress = self.g.run(query).data()
                answers += ress
            final_answer = self.answer_prettify(question_type, answers, sql_['entities'])
            if final_answer:
                answer.append(final_answer['answer'])
                nodes.extend(final_answer['nodes'])
                edges.extend(final_answer['edges'])

        entities = list(set(entities))  # 去重

        return {'answer': answer, 'entities': entities, 'relations': relations, 'KG': {'nodes': nodes, 'edges': edges}}

    '''根据对应的qustion_type，调用相应的回复模板'''

    def answer_prettify(self, question_type, answers, entities):
        final_answer = ''
        final_nodes = []
        final_edges = []
        upper_limit = 20
        #         '''
        #             已知头实体和尾实体，查询是否存在关系
        #         '''
        # is_symptom
        if question_type == 'is_symptom':
            # 即二者不存在关系
            if not answers:
                return {'answer': '{0}和{1}不存在“疾病与症状”之间的关系。'.format(entities[0], entities[1]), 'nodes': final_nodes,
                        'edges': final_edges}
            else:
                temp_answer = answers[0]
                nodes = dict([(temp_answer['labels(m)'][0], temp_answer['m.name']),
                              (temp_answer['labels(n)'][0], temp_answer['n.name'])])
                final_answer = '{0}会导致{1}，{1}是{0}的症状之一。'.format(nodes['Disease'], nodes['Symptom'])

                nodes_and_r = self.judgementDataConstruct(nodes, 'Disease', 'Symptom', temp_answer['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])


        elif question_type == 'is_belonging':
            if not answers:
                return {'answer': '{1}不是{0}的主治科室,{0}无法在{1}进行有效诊治。\n'.format(entities[0], entities[1]),
                        'nodes': final_nodes,
                        'edges': final_edges}
            else:
                temp_answer = answers[0]
                final_answer = '{1}是{0}的主治科室，{0}是在{1}进行挂号。'.format(temp_answer['m.name'], temp_answer['n.name'])

                nodes_and_r = self.judgementDataConstruct(
                    {'Disease': temp_answer['m.name'], 'Department': temp_answer['n.name']}, 'Disease', 'Department',
                    temp_answer['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'is_accompany_with':
            if not answers:
                return {'answer': '一般来说{0}和{1}并不存在并发症关系。\n'.format(entities[0], entities[1]), 'nodes': final_nodes,
                        'edges': final_edges}
            else:
                temp_answer = answers[0]
                temp_answer = str(temp_answer['p'])
                print(temp_answer)  # 形如(百日咳)<-[:acompany_with {}]-(小儿咳嗽)
                matchObj = re.match(r'[(](.*)[)].+[(](.*)[)]', temp_answer)
                # 判断方向
                if '>' in temp_answer:
                    # 正方向
                    start_node = matchObj.group(1)
                    end_node = matchObj.group(2)
                if '<' in temp_answer:
                    # 负方向
                    start_node = matchObj.group(2)
                    end_node = matchObj.group(1)

                final_answer = '{0}在某些情况下会导致{1}的并发。'.format(start_node, end_node)

                final_nodes.append({'id': self.count, 'name': end_node, 'category': self.categories['Disease']})
                temp1 = self.count
                self.count += 1
                final_nodes.append({'id': self.count, 'name': start_node, 'category': self.categories['Disease']})
                temp2 = self.count
                self.count += 1
                final_edges.append({'source': temp2, 'target': temp1, 'label': '并发症'})


        elif question_type == 'is_food':
            if not answers:
                return {'answer': '{1}并不存在对{0}的影响。'.format(entities[0], entities[1]), 'nodes': final_nodes,
                        'edges': final_edges}
            else:
                temp_answer = answers[0]
                if temp_answer['r.name'] == '推荐食谱':
                    final_answer = '患有{0}的病人推荐食用{1}。'.format(temp_answer['m.name'], temp_answer['n.name'])
                else:
                    final_answer = '患有{0}的病人{1}{2}。'.format(temp_answer['m.name'], temp_answer['r.name'],
                                                            temp_answer['n.name'])

                nodes_and_r = self.judgementDataConstruct(
                    {'Disease': temp_answer['m.name'], 'Food': temp_answer['n.name']}, 'Disease', 'Food',
                    temp_answer['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'is_need_check':
            if not answers:
                return {'answer': '{0}并不需要做{1}，通过此项检查无法有效检测出{0}。'.format(entities[0], entities[1]),
                        'nodes': final_nodes,
                        'edges': final_edges}
            else:
                temp_answer = answers[0]
                final_answer = '患有{0}的病人需要诊断检查{1}，通过{1}可以检查出{0}。'.format(temp_answer['m.name'], temp_answer['n.name'])

                nodes_and_r = self.judgementDataConstruct(
                    {'Disease': temp_answer['m.name'], 'Check': temp_answer['n.name']}, 'Disease', 'Check',
                    temp_answer['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'is_recommend_drug':
            if not answers:
                return {'answer': '{1}对{0}并无疗效。\n'.format(entities[0], entities[1]), 'nodes': final_nodes,
                        'edges': final_edges}
            else:
                if len(answers) == 2:
                    temp_answer_1 = answers[0]
                    temp_answer_2 = answers[1]
                    final_answer = '{1}是{0}的{2}以及{3}，患有{0}的病人在询问医师后可以使用{1}进行治疗。'.format(temp_answer_1['m.name'],
                                                                                        temp_answer_1['n.name'],
                                                                                        temp_answer_1['r.name'],
                                                                                        temp_answer_2['r.name'])
                    temp_answer = temp_answer_1
                elif len(answers) == 1:
                    temp_answer = answers[0]
                    final_answer = '{1}是{0}的{2}，患有{0}的病人在询问医师后可以使用{1}进行治疗。'.format(temp_answer['m.name'],
                                                                                   temp_answer['n.name'],
                                                                                   temp_answer['r.name'])

            nodes_and_r = self.judgementDataConstruct(
                {'Disease': temp_answer['m.name'], 'Drug': temp_answer['n.name']}, 'Disease', 'Drug',
                temp_answer['r.name'])
            final_nodes.extend(nodes_and_r[0])
            final_edges.extend(nodes_and_r[1])

        elif not answers:
            return ''

        if final_answer:
            # return {'answer': final_answer, 'KG': KG}
            return {'answer': final_answer, 'nodes': final_nodes, 'edges': final_edges}

        '''
            简单问句
        '''
        if question_type == 'disease_symptom':
            # 根据不同的疾病分组
            final_answer = ''
            for subject, items in groupby(answers, key=itemgetter('m.name')):
                desc = [i['n.name'] for i in items]
                final_answer += '{0}的症状包括：{1}。\n'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Disease', 'Symptom',
                                                           answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'symptom_disease':
            final_answer = ''
            for subject, items in groupby(answers, key=itemgetter('n.name')):
                desc = [i['m.name'] for i in items]
                final_answer += '可能具备症状{0}的疾病有：{1}。\n'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Symptom', 'Disease',
                                                           answers[0]['r.name'], direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])
            # desc = [i['m.name'] for i in answers]
            # subject = answers[0]['n.name']
            # final_answer = '可能具备症状{0}的疾病有：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))
            #
            # nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Symptom', 'Disease', answers[0]['r.name'], direction=False)
            # final_nodes.extend(nodes_and_r[0])
            # final_edges.extend(nodes_and_r[1])

        elif question_type == 'disease_cause':
            desc = [i['m.cause'] for i in answers]

            if len(desc) == 0 or desc[0] == '':
                subject = answers[0]['m.name']
                final_answer = '未找到{0}的成因。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0}可能的成因有：{1}。\n'.format(item['m.name'], item['m.cause'])

        elif question_type == 'disease_prevent':
            desc = [i['m.prevent'] for i in answers]

            if len(desc) == 0 or desc[0] == '':
                subject = answers[0]['m.name']
                final_answer = '未找到{0}的预防措施。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0}的预防措施：{1}。\n'.format(item['m.name'], item['m.prevent'])

        elif question_type == 'disease_lasttime':
            desc = [i['m.cure_lasttime'] for i in answers]

            if len(desc) == 0 or desc[0] == '':
                subject = answers[0]['m.name']
                final_answer = '未找到{0}的持续周期。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0}治疗可能持续的周期为：{1}。\n'.format(item['m.name'], item['m.cure_lasttime'])

        elif question_type == 'disease_cureway':
            desc = [';'.join(i['m.cure_way']) for i in answers]

            if len(desc) == 0 or desc[0] == '':
                subject = answers[0]['m.name']
                final_answer = '未找到{0}的治疗方案。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0}可以尝试如下治疗：{1}。\n'.format(item['m.name'], item['m.cure_way'])

        elif question_type == 'disease_cureprob':
            desc = [i['m.cured_prob'] for i in answers]
            if len(desc) == 0 or desc[0] == '':
                subject = answers[0]['m.name']
                final_answer = '未找到{0}的治愈概率。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0}治愈的概率为（仅供参考）：{1}。\n'.format(item['m.name'], item['m.cured_prob'])

        elif question_type == 'disease_easyget':
            desc = [i['m.easy_get'] for i in answers]
            subject = answers[0]['m.name']
            if len(desc) == 0 or desc[0] == '':
                final_answer = '未找到{0}的易感人群。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0}的易感人群有：{1}。\n'.format(item['m.name'], item['m.easy_get'])

        elif question_type == 'disease_desc':
            desc = [i['m.desc'] for i in answers]
            if len(desc) == 0 or desc[0] == '':
                subject = answers[0]['m.name']
                final_answer = '未找到{0}的相关描述。'.format(subject)
            else:
                final_answer = ''
                for item in answers:
                    final_answer += '{0},熟悉一下：{1}。\n'.format(item['m.name'], item['m.desc'])

        elif question_type == 'disease_acompany':
            final_answer = ''
            for subject in entities:
                desc = []
                for i in answers:
                    if i['m.name'] == subject:
                        desc.append(i['n.name'])
                    elif i['n.name'] == subject:
                        desc.append(i['m.name'])
                final_answer += '{0}的并发症包括：{1}。\n'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))
                # desc1 = [i['n.name'] for i in answers]
                # desc2 = [i['m.name'] for i in answers]
                # subject = answers[0]['m.name']
                # desc = [i for i in desc1 + desc2 if i != subject]
                # final_answer = '{0}的并发症包括：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Disease', 'Disease',
                                                           answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'disease_not_food':
            final_answer = ''
            for subject, items in groupby(answers, key=itemgetter('m.name')):
                desc = [i['n.name'] for i in items]
                # subject = answers[0]['m.name']
                final_answer += '{0}忌食的食物包括有：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Disease', 'Food',
                                                           answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'disease_do_food':
            final_answer = ''
            for subject in entities:
                do_desc = [i['n.name'] for i in answers if i['r.name'] == '宜吃' and i['m.name'] == subject]
                recommand_desc = [i['n.name'] for i in answers if i['r.name'] == '推荐食谱' and i['m.name'] == subject]
                # subject = answers[0]['m.name']
                final_answer += '{0}宜食的食物包括有：{1}；\n推荐食谱包括有：{2}。'.format(subject,
                                                                        ';'.join(list(set(do_desc))[:self.num_limit]),
                                                                        ';'.join(
                                                                            list(set(recommand_desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(do_desc + recommand_desc)), 'Disease',
                                                           'Food', answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'food_not_disease':
            final_answer = ''
            for subject in entities:
                desc = [i['m.name'] for i in answers if i['n.name'] == subject]
                # subject = answers[0]['n.name']
                if len(desc) != 0:
                    final_answer += '患有{0}的人最好不要吃{1}。'.format('；'.join(list(set(desc))[:self.num_limit]), subject)
                elif desc == []:
                    final_answer += '知识库中并未检索到{0}是某种疾病的忌食食物。'.format(subject)

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Food', 'Disease',
                                                           answers[0]['r.name'], direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'food_do_disease':
            final_answer = ''
            for subject in entities:
                desc = [i['m.name'] for i in answers if i['n.name'] == subject]
                # subject = answers[0]['n.name']
                if len(desc) != 0:
                    final_answer += '患有{0}的人建议多试试{1}。'.format('；'.join(list(set(desc))[:self.num_limit]), subject)
                elif desc == []:
                    final_answer += '知识库中并未检索到{0}是某种疾病的推荐食物。'.format(subject)

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Food', 'Disease',
                                                           answers[0]['r.name'], direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'disease_drug':
            final_answer = ''

            for subject in entities:
                desc = [i['n.name'] for i in answers if i['m.name'] == subject]
                # subject = answers[0]['m.name']
                final_answer += '{0}通常的使用的药品包括：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Disease', 'Drug',
                                                           answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'drug_disease':
            final_answer = ''

            for subject in entities:
                desc = [i['m.name'] for i in answers if i['n.name'] == subject]
                # subject = answers[0]['n.name']
                final_answer += '{0}主治的疾病有{1},可以试试。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Drug', 'Disease',
                                                           answers[0]['r.name'], direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'drug_producer':
            final_answer = ''

            for subject in entities:
                desc = [i['m.name'] for i in answers if i['n.name'] == subject]
                # subject = answers[0]['n.name']
                final_answer += '{0}在市面上的在售药品为：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Drug', 'Producer',
                                                           answers[0]['r.name'], direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'disease_check':
            final_answer = ''

            for subject in entities:
                desc = [i['n.name'] for i in answers if i['m.name'] == subject]
                # subject = answers[0]['m.name']
                final_answer += '{0}通常可以通过以下方式检查出来：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Disease', 'Check',
                                                           answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'check_disease':
            final_answer = ''

            for subject in entities:
                desc = [i['m.name'] for i in answers if i['n.name'] == subject]
                # subject = answers[0]['n.name']
                final_answer += '通常可以通过{0}检查出来的疾病有{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Check', 'Disease',
                                                           answers[0]['r.name'], direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'disease_department':
            final_answer = ''

            for subject in entities:
                desc = [i['n.name'] for i in answers if i['m.name'] == subject]
                # subject = answers[0]['m.name']
                final_answer += '得了{0}，应该看{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Disease', 'Department',
                                                           answers[0]['r.name'])
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        elif question_type == 'department_disease':
            final_answer = ''

            for subject in entities:
                desc = [i['m.name'] for i in answers if i['n.name'] == subject]
                # subject = answers[0]['n.name']
                if len(desc) != 0:
                    final_answer += '{0}主治的疾病包括：{1}。'.format(subject, '；'.join(list(set(desc))[:self.num_limit]))
                else:
                    final_answer += '知识库未检索到{0}的主治疾病。'.format(subject)

                nodes_and_r = self.simpleQuesDataConstruct(subject, list(set(desc)), 'Department', 'Disease',
                                                           answers[0]['r.name'],
                                                           direction=False)
                final_nodes.extend(nodes_and_r[0])
                final_edges.extend(nodes_and_r[1])

        '''
            复杂问句
        '''
        if question_type == 'intersection_symptom':
            final_answer = ''
            desc = [i['n1.name'] for i in answers]
            r_name = answers[0]['r1.name']
            if desc is None:
                final_answer = '{0}和{1}之间并没有相同的症状。'.format(entities[0], entities[1])
            else:
                # 存在相同症状
                final_answer = '{0}和{1}具有以下相同症状：{2}'.format(entities[0], entities[1],
                                                            '；'.join(list(set(desc))[:self.num_limit]))
                final_nodes.extend([
                    {'id': 0, 'name': entities[0], 'category': self.categories['Disease']},
                    {'id': 1, 'name': entities[1], 'category': self.categories['Disease']}
                ])
                count = 2
                for item in desc:
                    final_nodes.append({'id': count, 'name': item, 'category': self.categories['Symptom']})
                    final_edges.extend([
                        {'source': 0, 'target': count, 'label': r_name},
                        {'source': 1, 'target': count, 'label': r_name}
                    ])
                    count += 1

        elif question_type == 'intersection_disease':
            final_answer = ''
            desc = [i['m1.name'] for i in answers]
            r_name = answers[0]['r1.name']
            if desc is None:
                final_answer = '没有同时具备{0}和{1}这类症状的疾病。'.format(entities[0], entities[1])
            else:
                # 存在同时包含两种症状的疾病
                final_answer = '同时具有{0}和{1}这类症状的疾病有：{2}'.format(entities[0], entities[1],
                                                                '；'.join(list(set(desc))[:self.num_limit]))
                final_nodes.extend([
                    {'id': 0, 'name': entities[0], 'category': self.categories['Symptom']},
                    {'id': 1, 'name': entities[1], 'category': self.categories['Symptom']}
                ])
                count = 2
                for item in desc:
                    final_nodes.append({'id': count, 'name': item, 'category': self.categories['Disease']})
                    final_edges.extend([
                        {'source': count, 'target': 0, 'label': r_name},
                        {'source': count, 'target': 1, 'label': r_name}
                    ])
                    count += 1

        elif question_type == 'compare_disease_cured_prob':
            final_answer = ''
            desc_1 = answers[0]['m1.cured_prob']
            desc_2 = answers[0]['m2.cured_prob']
            pattern = "(\d+)%"
            l1 = re.findall(pattern, desc_1)
            l2 = re.findall(pattern, desc_2)
            if len(l1) == 0 or len(l2) == 0:
                # 无法比较
                final_answer += '{0}和{1}暂无法比较。'.format(entities[0], entities[1])
            else:
                l1_opposite = False
                l2_opposite = False
                opposite_words = ['死亡', '无']
                for word in opposite_words:
                    if word in desc_1:
                        l1_opposite = True
                    if word in desc_2:
                        l2_opposite = True
                is_greater = self.compare_prob(l1, l2, l1_opposite, l2_opposite)
                # 答案构造（无图谱）
                if is_greater == 1:
                    final_answer += '{0}的治愈率一般比{1}的治愈率更高。'.format(entities[0], entities[1])

                elif is_greater == -1:
                    final_answer += '{0}的治愈率一般比{1}的治愈率低。'.format(entities[0], entities[1])
                elif is_greater == 0:
                    final_answer += '{0}和[1}的治愈率相近。'.format(entities[0], entities[1])

            final_answer += '{0}的治愈率为：{2}；{1}的治愈率为{3}。'.format(entities[0], entities[1], desc_1, desc_2)

        elif question_type == 'compare_disease_lasttime':
            final_answer = ''
            desc_1 = answers[0]['m1.cure_lasttime']
            desc_2 = answers[0]['m2.cure_lasttime']
            if desc_1 is None or desc_2 is None:
                # 无法比较
                final_answer += '{0}和{1}暂无法比较。'.format(entities[0], entities[1])
            else:
                # 数字替换中文
                temp1 = self.text_to_num(desc_1)
                temp2 = self.text_to_num(desc_2)
                # 正则匹配
                pattern = "(\d+)(天|周|月|年|个月|星期|个星期)"
                l1 = re.findall(pattern, temp1)
                l2 = re.findall(pattern, temp2)

                lasttime_1 = self.time_translate(l1)
                lasttime_2 = self.time_translate(l2)

                if lasttime_1 > lasttime_2:
                    final_answer += '{0}的治疗周期相较{1}的一般会更长一些。'.format(entities[0], entities[1])
                elif lasttime_1 == lasttime_2:
                    final_answer += '{0}和{1}的治疗周期一般相近。'.format(entities[0], entities[1])
                elif lasttime_1 < lasttime_2:
                    final_answer += '{1}的治疗周期相较{0}的一般会更长一些。'.format(entities[0], entities[1])
            final_answer += '{0}的治疗周期为：{2}；{1}的治疗周期为{3}。'.format(entities[0], entities[1], temp1, temp2)

        result = {'answer': final_answer, 'nodes': final_nodes, 'edges': final_edges}
        return result

    def compare_prob(self, l1, l2, is_l1_opposite=False, is_l2_opposite=False):
        # 转换为可比较类型
        if is_l1_opposite:
            t1 = [100 - int(x) for x in l1]
        else:
            t1 = [int(x) for x in l1]

        if is_l2_opposite:
            t2 = [100 - int(x) for x in l2]
        else:
            t2 = [int(x) for x in l2]

        t1.sort()
        t2.sort()
        prob_1 = t1[-1]
        prob_2 = t2[-1]
        if prob_1 > prob_2:
            return 1
        elif prob_1 == prob_2:
            return 0
        else:
            return -1

    def text_to_num(self, sen):
        # 将文本中的“一二三”替换为可比较数字“123”
        mapper = {'一': '1', '二': '2', '三': '3', '四': '4', '五': '5', '六': '6', '七': '7', '八': '8', '九': '9', '十': '10', }
        for word in sen:
            if word in mapper.keys():
                sen = sen.replace(word, mapper[word])

        return sen

    def time_translate(self, times):
        # times形如[('3', '周'), ('15', '天')]
        mapper = {'天': 1, '周': 7, '月': 30, '年': 365, '星期': 7, '个月': 30, '个星期': 7}
        res = []
        for t in times:
            time = int(t[0]) * mapper[t[1]]
            res.append(time)

        res.sort()
        return res[-1]

    def judgementDataConstruct(self, nodes, label_1, label_2, r_name):
        '''
        @param nodes: {类型名：[具体实体，具体实体]}
        @param label_1: 主体实体类别
        @param label_2: 目标实体类别
        @return: (temp_nodes, temp_relations)
        '''
        temp_nodes = []
        temp_relations = []
        temp_nodes.append({'id': self.count, 'name': nodes[label_1], 'category': self.categories[label_1]})
        self.count += 1
        temp_nodes.append({'id': self.count, 'name': nodes[label_2], 'category': self.categories[label_2]})
        self.count += 1
        temp_relations.append({'source': temp_nodes[0]['id'], 'target': temp_nodes[1]['id'], 'label': r_name})

        return temp_nodes, temp_relations

    def simpleQuesDataConstruct(self, parent, children, label_1, label_2, r_name, direction=True, upper_bounce=30):
        '''
        @param children: 子节点
        @param parent: 父节点
        @param label_1: 父节点标签
        @param label_2: 子节点标签
        @param direction: 连线方向
        @param upper_bounce: 子节点上限，默认为30
        @return: (temp_nodes, temp_relations)
        '''
        temp_nodes = []
        temp_relations = []
        if len(children) > upper_bounce:
            children = children[:upper_bounce]

        temp_nodes.append({'id': self.count, 'name': parent, 'category': self.categories[label_1]})
        start = self.count
        self.count += 1

        for node in children:
            temp_nodes.append({'id': self.count, 'name': node, 'category': self.categories[label_2]})
            self.count += 1
        end = self.count

        if direction:
            for i in range(start + 1, end):
                temp_relations.append({'source': start, 'target': i, 'label': r_name})
        else:
            for i in range(start + 1, end):
                temp_relations.append({'source': i, 'target': start, 'label': r_name})

        return temp_nodes, temp_relations


if __name__ == '__main__':
    searcher = AnswerSearcher()
    res_sql = [{
        'question_type': 'disease_cureway',
        'entities': ['肾虚'],
        'sql': ["MATCH (m:Disease) where m.name = '肾虚' return m.name, m.cure_way"]
    }]

    final_answers = searcher.search_main(res_sql)
    print(final_answers)

